#include <iostream>
#include <assert.h>
using namespace std;

enum Color{
  RED,
  BLACK
};

template <class K, class V>
struct RBTreeNode{
  RBTreeNode<K,V> *_left;
  RBTreeNode<K,V> *_right;
  RBTreeNode<K,V> *_parent;
  pair<K,V> _kv;
  Color _color;

  RBTreeNode(const pair<K,V> &kv=pair<K,V>(), Color color = RED)
    :_left(nullptr),
    _right(nullptr),
    _parent(nullptr),
    _kv(kv),
    _color(color)
  {}
};

template <class K, class V>
class RBTree{
  typedef RBTreeNode<K,V> Node;
  Node *_phead;

public:
  RBTree(){
    _phead = new Node;
    _phead->_left = _phead;
    _phead->_right = _phead;
  }

  Node*& GetRoot(){
    return _phead->_parent;
  }

  Node* LeftMost();
  Node* RightMost();
  bool Insert(const pair<K,V> &kv);
  Node* Find(const K &k);

  void Inorder(){
    _Inorder(GetRoot());
    cout << endl;
  }

  bool IsValidRBTree();

private:
  void RotateL(Node *parent);
  void RotateR(Node *parent);
  void _Inorder(Node *root);
  bool _IsValidRBTree(Node *root, int blacknum, int &benchmark);
};

template <class K, class V>
typename RBTree<K,V>::Node* RBTree<K,V>::LeftMost(){
  Node *root = GetRoot();
  if(root == nullptr)
    return _phead;
  else{
    Node *left = root;
    while(left->_left!=nullptr)
    {
      left = left->_left;
    }
    return left;
  }
}

template <class K, class V>
typename RBTree<K,V>::Node* RBTree<K,V>::RightMost(){
  Node *root = GetRoot();
  if(root == nullptr)
    return _phead;
  else{
    Node *right = root;
    while(right->_right!=nullptr)
    {
      right = right->_right;
    }
    return right;
  }
}

template <class K, class V>
bool RBTree<K,V>::IsValidRBTree(){
    Node *root = GetRoot();
  //空树也是红黑树
  if(root == nullptr) return true;
  //检查性质2:
  if(root->_color != BLACK)
  {
    cout << "违反性质2：根节点不为黑色！" << endl;
    return false;
  }
  //检查性质3,4：
  int benchmark = 0;
  return _IsValidRBTree(root, 0, benchmark);
}

template <class K, class V>
bool RBTree<K,V>::_IsValidRBTree(Node *root, int blacknum, int &benchmark){
  if(root == nullptr)
  {
    if(benchmark == 0) //表示第一条路径遍历完
    {
      benchmark = blacknum; //记录第一条路径的黑色节点个数
      return true;
    }
    else{
      if(blacknum != benchmark) //如果其他路径的blacknum与第一条路径不同，说明违反性质4
      {
        cout << "违反性质4：从任意节点到每个叶子节点的所有路径都包含相同数目的黑色节点！" << endl;
        return false;
      }
      else{
        return true;
      }
    }
  }
    
  //检查性质3：
  if(root->_color == RED && root->_parent->_color == RED)
  {
    cout << "违反性质3：路径上有两个连续的红色节点！" << endl;
    return false;
  }

  if(root->_color == BLACK)
  {
    ++blacknum; 
  }
  return _IsValidRBTree(root->_left, blacknum, benchmark)
      && _IsValidRBTree(root->_right, blacknum, benchmark);
}

template <class K, class V>
void RBTree<K,V>::_Inorder(Node *root){
  if(root == nullptr) return; 
  _Inorder(root->_left);
  cout << root->_kv.first << ":" << root->_kv.second << " ";
  _Inorder(root->_right);
}

template <class K, class V>
bool RBTree<K,V>::Insert(const pair<K,V> &kv)
{
  Node* &root = GetRoot();
  if(root == nullptr)
  {
    root = new Node(kv, BLACK);
    root->_parent = _phead;
    _phead->_left = root;
    _phead->_right = root;
    return true;
  }

  Node *cur = root;
  Node *parent = nullptr;
  while(cur != nullptr)
  {
    if(kv.first > cur->_kv.first)
    {
      parent = cur;
      cur = cur->_right;
    }
    else if(kv.first < cur->_kv.first)
    {
      parent  = cur;
      cur = cur->_left;
    }
    else{
      return false;
    }
  }
  cur = new Node(kv,RED);
  if(kv.first > parent->_kv.first)
  {
    parent->_right = cur;
  }
  else{
    parent->_left = cur;
  }
  cur->_parent = parent;
  
  while(parent != _phead && parent->_color == RED) //上一次循环中grandparent 为根节点，此次循环parent == nullptr
  {
    Node *grandparent = parent->_parent;
    assert(grandparent != nullptr);
    assert(grandparent->_color == BLACK);

    Node *uncle = grandparent->_left;
    if(parent == grandparent->_left)
      uncle = grandparent->_right;

    if(uncle != nullptr && uncle->_color == RED) //情况一：uncle存在且为红
    {
      parent->_color = uncle->_color = BLACK;
      grandparent->_color = RED;
      cur = grandparent;
      parent = cur->_parent;
    }
    else //情况二、三：uncle不存在或uncle存在且为黑
    {
      if(parent == grandparent->_left)
      {
        if(cur == parent->_left) //左左
        {
          RotateR(grandparent);
          parent->_color = BLACK;
          grandparent->_color = RED;
        }
        else{ //左右
          RotateL(parent);
          RotateR(grandparent);
          cur->_color = BLACK;
          grandparent->_color = RED;
        }
      }
      else{
        if(cur == parent->_right) //右右
        {
          RotateL(grandparent);
          parent->_color = BLACK;
          grandparent->_color = RED;
        }
        else{ //右左
          RotateR(parent);
          RotateL(grandparent);
          cur->_color = BLACK;
          grandparent->_color = RED;
        }
      }
      break;
    } //end of else
  } //end of while
  
  if(parent == _phead)
    root->_color = BLACK;
  _phead->_left = LeftMost();
  _phead->_right = RightMost();
  return true;
}

template <class K, class V>
void RBTree<K,V>::RotateL(Node *parent){
  Node *subR = parent->_right;
  Node *subRL = subR->_left;
  Node *ppNode = parent->_parent;
  Node* &root = GetRoot();

  parent->_right = subRL;
  if(subRL != nullptr)
  {
    subRL->_parent = parent;
  }

  subR->_left = parent;
  parent->_parent = subR;
 
  if(root == parent)
  {
    root = subR;
    subR->_parent = _phead;
  }
  else{
    if(ppNode->_left == parent)
    {
      ppNode->_left = subR;
    }
    else{
      ppNode->_right = subR;
    }
  }
  subR->_parent = ppNode;
}

template <class K, class V>
void RBTree<K,V>::RotateR(Node *parent){
  Node *subL = parent->_left;
  Node *subLR = subL->_right;
  Node *ppNode = parent->_parent;
  Node* &root = GetRoot();

  parent->_parent = subL;
  subL->_right = parent;

  parent->_left = subLR;
  if(subLR != nullptr)
  subLR->_parent = parent;

  if(root == parent)
  {
    root = subL;
    subL->_parent = _phead;
  }
  else{
    if(ppNode->_left == parent)
    {
      ppNode->_left = subL;
    }
    else{
      ppNode->_right = subL;
    }
  }
  subL->_parent = ppNode;
}

template <class K, class V>
typename RBTree<K,V>::Node* RBTree<K,V>::Find(const K &k){
  Node *root = GetRoot();
  if(root == nullptr)
    return nullptr;
  Node *cur = root;
  while(cur != nullptr)
  {
    if(k > cur->_kv.first)
    {
      cur = cur->_right;
    }
    else if(k < cur->_kv.first)
    {
      cur = cur->_left;
    }
    else{
      return cur;
    }
  }
  return nullptr;
}
